# This code is part of cqlib.
#
# Copyright (C) 2024 China Telecom Quantum Group, QuantumCTek Co., Ltd.,
# Center for Excellence in Quantum Information and Quantum Physics.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

from typing import Optional, Union, List, Dict
import numpy as np
from matplotlib import pyplot as plt

VALID_SORTS = {"asc", "desc", "value", "value_desc", "hamming"}

def hamming_distance_bits(a: str, b: str) -> int:
    """
    Compute the Hamming distance between two equal-length bitstrings.

    Args:
        a (str): First bitstring (e.g., "0011").
        b (str): Second bitstring with the same length as ``a``.

    Returns:
        int: Number of positions where the bits differ.

    Raises:
        ValueError: If ``a`` and ``b`` do not have the same length.
    """
    if len(a) != len(b):
        raise ValueError("Hamming sort needs equal-length bitstrings.")
    return sum(x != y for x, y in zip(a, b))


def _keep_topk_with_rest(exe: Dict[str, float], k: int) -> Dict[str, float]:
    """
    Keep only the top-K items and aggregate the remainder into ``"rest"``.

    Args:
        exe (dict[str, float]): Label-to-value mapping.
        k (int): Number of largest items to keep. If ``k`` is None/<=0 or
            ``k >= len(exe)``, returns a shallow copy of ``exe`` unchanged.

    Returns:
        dict[str, float]: New mapping with top-K items and (if applicable) a
            ``"rest"`` entry equal to the sum of dropped values.
    """
    if k is None or k <= 0 or k >= len(exe):
        return dict(exe)
    sorted_items = sorted(exe.items(), key=lambda kv: kv[1])
    rest = sum(v for _, v in sorted_items[:-k])
    kept = dict(sorted_items[-k:])
    kept["rest"] = rest
    return kept


def _align_labels(dicts: List[Dict[str, float]]) -> List[Dict[str, float]]:
    """
    Align key sets across multiple dictionaries, filling missing keys with 0.0.

    Args:
        dicts (list[dict[str, float]]): Dictionaries that may have different keys.

    Returns:
        list[dict[str, float]]: New dictionaries sharing the union of all keys,
            suitable for grouped bar alignment.
    """
    all_keys = set().union(*(d.keys() for d in dicts))
    base = {k: 0.0 for k in all_keys}
    out = []
    for d in dicts:
        nd = base.copy()
        nd.update(d)
        out.append(nd)
    return out


def draw_histogram(
    histogram_data: Optional[Union[Dict[str, float], List[Dict[str, float]]]],
    title: Optional[str] = None,
    bar_value_show: bool = True,
    file_name: Optional[str] = None,
    axis_label: Optional[List[str]] = None,
    interactive: bool = False,
    fig_size: tuple = (6.4, 4.8),
    legend: Optional[List[str]] = None,
    *,
    sort: str = "asc",
    target_string: Optional[str] = None,
    number_to_keep: Optional[int] = None,
    color: Optional[Union[str, List[str]]] = None,
    ax: Optional[plt.Axes] = None,
    return_fig: bool = True,
    hide_zero: bool = True,
):
    """
    Draw a grouped histogram from measurement results (counts, probabilities, or quasi-probabilities).

    Values are plotted **as provided** (no normalization). Supports multi-dataset grouped bars,
    x-axis label sorting, per-dataset Top-K with 'rest' aggregation, optional hiding of zero-valued
    items, custom colors/legend, file saving, and notebook-friendly rendering.

    Args:
        histogram_data (dict[str, float] | list[dict[str, float]] | None):
            Mapping(s) from label (e.g., bitstring) to numeric value. A single dict
            produces one series; a list of dicts produces grouped bars.
        title (str | None):
            Figure title displayed above the chart.
        bar_value_show (bool):
            Whether to render numeric value labels above each bar. Defaults to True.
        file_name (str | None):
            Optional output path. If provided, the figure is saved to this path
            (no automatic suffix is appended).
        axis_label (list[str] | None):
            Axis titles as ``[x_label, y_label]``. If omitted or length < 2, the
            y-axis title defaults to ``"Value"``.
        interactive (bool):
            If True, immediately shows the figure via ``plt.show()`` and closes it
            to avoid duplicate rendering in notebooks.
        fig_size (tuple[float, float]):
            Figure size in inches, as ``(width, height)``. Defaults to ``(6.4, 4.8)``.
        legend (list[str] | None):
            Legend labels for each dataset. Length must equal the number of datasets.
        sort ({"asc", "desc", "value", "value_desc", "hamming"}):
            Strategy for ordering x-axis labels:
              - ``"asc"`` / ``"desc"``: lexicographic order of labels.
              - ``"value"`` / ``"value_desc"``: ascending/descending by the **maximum**
                value across datasets for each label.
              - ``"hamming"``: ascending by Hamming distance to ``target_string``.
            Defaults to ``"asc"``.
        target_string (str | None):
            Bitstring of the same length as labels, required when ``sort="hamming"``.
        number_to_keep (int | None):
            Keep the top-K labels **per dataset** and aggregate the remainder into a
            special label ``"rest"``. After this step, label sets are aligned across datasets.
        color (str | list[str] | None):
            A single color or a list of colors for datasets. Defaults to Matplotlib's rc cycle.
        ax (matplotlib.axes.Axes | None):
            External axes to draw on. When supplied, the function draws onto this axes
            and returns ``None``.
        return_fig (bool):
            If True and ``ax`` is None, returns the created ``Figure``. In notebook inline
            backends, returning a ``Figure`` may cause it to render if it is the cell's
            last expression. Defaults to False.
        hide_zero (bool):
            If True, removes labels that are zero across **all** datasets and skips
            drawing individual zero-valued bars within remaining labels. Defaults to True.

    Returns:
        matplotlib.figure.Figure | None:
            The created figure when ``ax`` is None and ``return_fig`` is True; otherwise ``None``.

    Raises:
        ValueError:
            If ``histogram_data`` is None; if ``sort`` is invalid; if ``sort="hamming"``
            without ``target_string``; or if ``legend`` is provided with a length that
            does not match the number of datasets.

    Notes:
        - X-axis tick labels are the (possibly aggregated) labels, centered per group and
          rotated 70° for readability.
        - Negative values (e.g., quasi-probabilities) are supported; the y-axis range
          is expanded accordingly.
        - To avoid duplicate images in Jupyter/IPython, prefer ``interactive=True`` or
          set ``return_fig=False`` and avoid leaving the returned figure as the last
          expression of a cell.
    """
    if histogram_data is None:
        raise ValueError("histogram_data is None.")
    if sort not in VALID_SORTS:
        raise ValueError(f"sort must be one of {sorted(VALID_SORTS)}")
    if sort == "hamming" and target_string is None:
        raise ValueError("sort='hamming' requires target_string.")

    # Normalize input to a list of dicts for unified processing.
    if isinstance(histogram_data, dict):
        data_list = [dict(histogram_data)]
    else:
        data_list = [dict(d) for d in histogram_data]

    # Optional per-dataset Top-K + key alignment (ensures grouped bars share the same labels).
    if number_to_keep is not None:
        data_list = [_keep_topk_with_rest(d, number_to_keep) for d in data_list]
        data_list = _align_labels(data_list)

    # Gather and sort labels for the x-axis.
    all_labels = sorted(set().union(*(d.keys() for d in data_list)))
    if sort == "hamming":
        labels_wo_rest = [x for x in all_labels if x != "rest"]
        labels_wo_rest.sort(key=lambda s: hamming_distance_bits(s, target_string))
        labels = labels_wo_rest + (["rest"] if "rest" in all_labels else [])
    elif sort in {"value", "value_desc"}:
        # Use the maximum value across datasets per label to derive a global importance order.
        peak = {k: max(d.get(k, 0.0) for d in data_list) for k in all_labels}
        labels = sorted(all_labels, key=lambda k: peak[k], reverse=(sort == "value_desc"))
    else:
        labels = list(all_labels)
        if sort == "desc":
            labels.reverse()

    # Resolve colors.
    if color is None:
        color_cycle = plt.rcParams["axes.prop_cycle"].by_key().get("color", ["C0"])
    elif isinstance(color, str):
        color_cycle = [color]
    else:
        color_cycle = list(color)

    # Create/resolve axes and figure.
    created_fig = False
    if ax is None:
        fig, ax = plt.subplots(figsize=fig_size)
        created_fig = True
    else:
        fig = ax.figure

    # Construct value arrays following `labels` order (no normalization is applied).
    values_list = [np.array([float(d.get(k, 0.0)) for k in labels], dtype=float) for d in data_list]

    # Optionally drop labels that are zero across **all** datasets, and skip zero bars.
    if hide_zero:
        if values_list:
            stacked = np.vstack(values_list)  # shape: (n_sets, n_labels)
            keep_cols = np.any(np.abs(stacked) > 0, axis=0)
        else:
            keep_cols = np.array([], dtype=bool)
        labels = [lbl for lbl, keep in zip(labels, keep_cols) if keep]
        values_list = [vals[keep_cols] for vals in values_list]

    # Compute bar placement for grouped bars (centered on integer positions).
    n_sets = len(values_list)
    n_labels = len(labels)
    x_centers = np.arange(n_labels)
    group_width = 1.0
    bar_width = group_width / (n_sets + 1 if n_sets > 0 else 2)

    # Draw each dataset as a group-offset series of bars.
    for i, vals in enumerate(values_list):
        x_pos_all = x_centers + (i - (n_sets - 1) / 2) * bar_width
        if hide_zero:
            nz_idx = np.where(np.abs(vals) > 0)[0]
            x_pos = x_pos_all[nz_idx]
            vals_to_plot = vals[nz_idx]
        else:
            x_pos = x_pos_all
            vals_to_plot = vals

        lbl = legend[i] if legend and i < len(legend) else None
        rects = ax.bar(
            x_pos,
            vals_to_plot,
            width=bar_width,
            label=lbl,
            zorder=2,
            color=color_cycle[i % len(color_cycle)],
            align="center",
        )

        # Optional value labels above bars.
        if bar_value_show and len(vals_to_plot) > 0:
            dy = max(0.02, 0.01 * float(np.max(np.abs(vals_to_plot))))
            for r, v in zip(rects, vals_to_plot):
                txt = f"{v:.0f}" if abs(v) >= 1 else f"{v:.3f}"
                ax.text(
                    r.get_x() + r.get_width() / 2,
                    v + dy,
                    txt,
                    ha="center",
                    va="bottom",
                    zorder=3,
                )

    # Configure x-axis ticks and labels at group centers; rotate for readability.
    ax.set_xticks(x_centers)
    ax.set_xticklabels(labels, rotation=70, ha="right", rotation_mode="anchor")

    # Y-axis range: allow negative values (e.g., quasi-probabilities).
    if values_list and any(len(v) for v in values_list):
        all_vals = np.concatenate(values_list)
        if hide_zero:
            all_vals = all_vals[np.abs(all_vals) > 0]
        if all_vals.size == 0:
            vmin, vmax = 0.0, 1.0
        else:
            vmin, vmax = float(np.min(all_vals)), float(np.max(all_vals))
    else:
        vmin, vmax = 0.0, 1.0

    y_min = min(0.0, vmin * 1.1)
    y_max = max(1.0 if vmax == 0 else vmax * 1.1, 1e-2)
    if y_max - y_min < 1e-12:
        y_max = y_min + 1.0
    ax.set_ylim(y_min, y_max)

    # Axis titles.
    if axis_label and len(axis_label) >= 1:
        ax.set_xlabel(axis_label[0])
    ax.set_ylabel(axis_label[1] if axis_label and len(axis_label) >= 2 else "Value")

    # Grid, legend, and title.
    ax.grid(which="major", axis="y", linestyle="--", zorder=0)
    if legend:
        if len(legend) != n_sets:
            raise ValueError(f"legend长度为{len(legend)}，但数据集数量为{n_sets}，不匹配。")
        ax.legend(loc="upper left", bbox_to_anchor=(1.01, 1.0), borderaxespad=0, frameon=True)
    if title:
        ax.set_title(title)

    # Save if requested.
    if file_name:
        try:
            fig.tight_layout()
        except Exception:
            pass
        fig.savefig(file_name, dpi=200)

    # Jupyter inline double-render guard.
    backend = plt.get_backend().lower()
    is_inline = ("inline" in backend) or ("nbagg" in backend)

    if interactive:
        plt.show()
        plt.close(fig)
        return None

    if is_inline and not return_fig and ax is None:
        try:
            fig.tight_layout()
        except Exception:
            pass
        plt.close(fig)
        return None

    if return_fig:
        try:
            fig.tight_layout()
        except Exception:
            pass
        return fig

    # External Axes provided → nothing to return.
    return None